import os
import csv
import numpy as np
import sys
from typing import Dict, List, Optional, Tuple


class StatsUtil:
    """
    统计文件处理工具类，负责与统计文件交互，提供各种统计信息的加载和计算功能
    """
    
    # 静态缓存，避免重复加载统计文件
    _stats_cache = {}
    
    # 睡眠阶段映射：N3和N4在某些数据集可能合并
    STAGE_MAPPING = {
        'wake_count': 0,
        'n1_count': 1,
        'n2_count': 2,
        'n3_count': 3,
        'n4_count': 3,  # N4阶段映射到N3
        'rem_count': 4
    }
    
    # CSV文件的列索引映射（基于文件格式）
    CSV_COLUMN_MAPPING = {
        0: 'filename',
        1: 'total_stages',  # 第一列是文件名，第二列是总阶段数
        2: 'wake_count',
        3: 'n1_count', 
        4: 'n2_count',
        5: 'n3_count',
        6: 'rem_count',
        7: 'sampling_rate',  # 采样率 - 修正索引
        8: 'data_mean',      # 数据均值 - 修正索引
        9: 'data_std'        # 数据标准差 - 修正索引
    }
    
    @classmethod
    def convert_to_absolute_path(cls, relative_path: str) -> str:
        """
        将相对路径转换为绝对路径
        
        Args:
            relative_path: 相对路径
            
        Returns:
            绝对路径
        """
        if os.path.isabs(relative_path):
            return relative_path
        
        # 尝试从项目根目录开始解析相对路径
        # 查找项目根目录（包含training文件夹的目录）
        current_dir = os.path.dirname(os.path.abspath(__file__))
        while True:
            parent_dir = os.path.dirname(current_dir)
            if parent_dir == current_dir:  # 到达文件系统根目录
                break
            
            if os.path.exists(os.path.join(parent_dir, 'training')):
                # 构建基于项目根目录的绝对路径
                absolute_path = os.path.join(parent_dir, relative_path)
                return absolute_path
                
            current_dir = parent_dir
        
        # 如果找不到项目根目录，使用当前工作目录
        return os.path.abspath(relative_path)
    
    @classmethod
    def load_statistics_file(cls, stats_file: Optional[str]) -> Dict[str, Dict]:
        """
        加载统计文件，返回文件统计信息字典
        使用缓存机制避免重复加载
        
        Args:
            stats_file: 统计文件路径
            
        Returns:
            字典，键为文件名，值为包含各阶段统计信息的字典
        """
        if stats_file is None:
            print("统计文件路径为None，无法加载")
            return {}
        
        try:
            # 转换为绝对路径
            absolute_path = cls.convert_to_absolute_path(stats_file)
            
            # 检查缓存
            if absolute_path in cls._stats_cache:
                return cls._stats_cache[absolute_path]
                
            if not os.path.exists(absolute_path):
                print(f"统计文件不存在: {absolute_path}")
                return {}
            
            file_stats = {}
            with open(absolute_path, 'r', encoding='utf-8') as f:
                reader = csv.reader(f)
                for row in reader:
                    if len(row) < 8:  # 确保有足够的列
                        continue
                    
                    filename = row[0].strip()
                    
                    # 构建文件的统计信息字典
                    stats = {}
                    for col_idx, col_name in cls.CSV_COLUMN_MAPPING.items():
                        if col_idx < len(row):
                            # 特殊处理data_mean和data_std为浮点数
                            if col_name in ['data_mean', 'data_std', 'sampling_rate']:
                                try:
                                    stats[col_name] = float(row[col_idx])
                                except ValueError:
                                    stats[col_name] = 0.0 if col_name != 'sampling_rate' else 0
                            else:
                                # 其他字段为整数
                                try:
                                    stats[col_name] = int(row[col_idx])
                                except ValueError:
                                    stats[col_name] = 0
                    
                    # 计算总阶段数（如果没有提供）
                    if 'total_stages' not in stats or stats['total_stages'] == 0:
                        stage_counts = [stats.get(stage, 0) for stage in ['wake_count', 'n1_count', 'n2_count', 'n3_count', 'rem_count']]
                        stats['total_stages'] = sum(stage_counts)
                    
                    file_stats[filename] = stats
            
            print(f"成功加载统计文件，包含 {len(file_stats)} 个文件的统计信息")
            # 缓存结果
            cls._stats_cache[absolute_path] = file_stats
            return file_stats
            
        except Exception as e:
            print(f"加载统计文件失败: {str(e)}")
            return {}
    
    @classmethod
    def get_file_stages_count(cls, stats_file: Optional[str], filenames: List[str]) -> Dict[str, int]:
        """
        获取指定文件列表的阶段数量
        
        Args:
            stats_file: 统计文件路径
            filenames: 文件列表
            
        Returns:
            字典，键为文件名，值为阶段数量
        """
        print(f"从统计文件获取 {len(filenames)} 个文件的阶段数量")
        
        file_stages_map = {}
        missing_files = []
        
        # 为每个文件查找阶段数量
        for filename in filenames:
            # 使用get_file_stats方法获取文件统计信息
            # 由于get_file_stats需要完整路径，这里使用文件名作为路径参数
            file_stats = cls.get_file_stats(filename, stats_file)
            
            if file_stats and 'stages_count' in file_stats:
                file_stages_map[filename] = file_stats['stages_count']
            else:
                missing_files.append(filename)
        
        if missing_files:
            print(f"统计文件中缺少 {len(missing_files)} 个文件的信息")
        
        return file_stages_map
    
    @classmethod
    def calculate_class_weights(cls, stats_file: Optional[str], filenames: Optional[List[str]] = None) -> Optional[List[float]]:
        """
        从统计文件计算类别权重
        
        Args:
            stats_file: 统计文件路径
            filenames: 可选，特定的文件列表，如果为None则使用所有文件
            
        Returns:
            类别权重列表，如果无法计算则返回None
        """
        # 加载统计文件获取所有文件列表（用于后续处理）
        all_file_stats = cls.load_statistics_file(stats_file)
        
        if not all_file_stats:
            print("统计文件中没有有效的类别统计数据")
            return None
        
        # 确定要使用的文件集合
        if filenames is not None:
            target_files = filenames
        else:
            # 如果没有指定文件名列表，使用统计文件中的所有文件
            target_files = list(all_file_stats.keys())
        
        # 计算各类别总数
        class_counts = [0] * 5  # 5个睡眠阶段类别
        
        for filename in target_files:
            # 使用get_file_stats方法获取文件统计信息
            file_stats = cls.get_file_stats(filename, stats_file)
            
            if file_stats:
                # 统计各类别数量（处理N3和N4合并的情况）
                for stage_name, class_idx in cls.STAGE_MAPPING.items():
                    if class_idx < len(class_counts):  # 确保索引有效
                        class_counts[class_idx] += file_stats.get(stage_name, 0)
        
        # 检查是否有有效数据
        total_samples = sum(class_counts)
        if total_samples == 0:
            print("计算得到的总样本数为0，无法计算类别权重")
            return None
        
        # 计算权重：使用倒数频率并归一化
        weights = [total_samples / max(count, 1) for count in class_counts]
        
        # 归一化权重
        weights_sum = sum(weights)
        if weights_sum > 0:
            weights = [w / weights_sum * len(weights) for w in weights]
        
        print(f"成功计算类别权重: {weights}")
        return weights
    
    @classmethod
    def get_dataset_class_distribution(cls, stats_file: Optional[str], filenames: List[str]) -> Dict[str, int]:
        """
        获取数据集的类别分布
        
        Args:
            stats_file: 统计文件路径
            filenames: 文件列表
            
        Returns:
            类别分布字典
        """
        # 类别分布
        class_dist = {
            'wake': 0,
            'n1': 0,
            'n2': 0,
            'n3': 0,
            'rem': 0
        }
        
        for filename in filenames:
            # 使用get_file_stats方法获取文件统计信息
            file_stats = cls.get_file_stats(filename, stats_file)
            
            if file_stats:
                class_dist['wake'] += file_stats.get('wake_count', 0)
                class_dist['n1'] += file_stats.get('n1_count', 0)
                class_dist['n2'] += file_stats.get('n2_count', 0)
                # N3和N4合并
                class_dist['n3'] += file_stats.get('n3_count', 0) + file_stats.get('n4_count', 0)
                class_dist['rem'] += file_stats.get('rem_count', 0)
        
        return class_dist
    
    @classmethod
    def validate_statistics_file(cls, stats_file: str) -> Tuple[bool, str]:
        """
        验证统计文件的有效性
        
        Args:
            stats_file: 统计文件路径
            
        Returns:
            (是否有效, 消息)
        """
        try:
            absolute_path = cls.convert_to_absolute_path(stats_file)
            
            if not os.path.exists(absolute_path):
                return False, f"文件不存在: {absolute_path}"
            
            with open(absolute_path, 'r', encoding='utf-8') as f:
                reader = csv.reader(f)
                line_count = 0
                valid_lines = 0
                
                for row in reader:
                    line_count += 1
                    if len(row) >= 8:  # 至少需要8列
                        try:
                            # 尝试转换数值列
                            for i in range(1, len(row)):
                                int(row[i])
                            valid_lines += 1
                        except ValueError:
                            continue
                
                if line_count == 0:
                    return False, "统计文件为空"
                
                if valid_lines == 0:
                    return False, "统计文件中没有有效的数据行"
                
                return True, f"统计文件有效，包含 {valid_lines}/{line_count} 行有效数据"
                
        except Exception as e:
            return False, f"验证统计文件失败: {str(e)}"
            
    @classmethod
    def get_file_stats(cls, file_path, stats_file):
        """
        从统计文件中获取特定文件的统计信息
        
        Args:
            file_path: 文件路径
            stats_file: 统计文件路径
            
        Returns:
            字典，包含文件的统计信息，特别是'stages_count'键
        """
        if stats_file is None:
            print("统计文件路径为None，无法加载文件统计信息")
            return None
        
        try:
            # 获取文件名
            filename = os.path.basename(file_path)
            
            # 加载统计文件
            file_stats = cls.load_statistics_file(stats_file)
            
            if filename in file_stats:
                # 提取文件统计信息
                stats = file_stats[filename]
                
                # 构建返回字典，确保包含'stages_count'键
                result = {
                    'stages_count': stats.get('total_stages', 0)
                }
                
                # 添加其他可能的统计信息
                for stage_name in ['wake_count', 'n1_count', 'n2_count', 'n3_count', 'rem_count']:
                    if stage_name in stats:
                        result[stage_name] = stats[stage_name]
                
                return result
            else:
                print(f"统计文件中未找到文件 {filename} 的信息")
                return None
                
        except Exception as e:
            print(f"加载文件统计信息失败: {str(e)}")
            return None
    
    @classmethod
    def get_normalization_stats_from_stats(cls, stats_file):
        """
        从统计文件中获取用于标准化的均值和标准差
        
        Args:
            stats_file: 统计文件路径
            
        Returns:
            字典，包含'mean'和'std'键，分别对应均值和标准差列表
            如果无法从统计文件获取，则返回空字典
        """
        if stats_file is None:
            print("统计文件路径为None，无法加载标准化统计信息")
            return {}
        
        try:
            # 转换为绝对路径
            absolute_path = cls.convert_to_absolute_path(stats_file)
            if not os.path.exists(absolute_path):
                print(f"统计文件不存在: {absolute_path}")
                return {}
            
            # 尝试从统计文件中读取标准化信息
            file_stats = cls.load_statistics_file(stats_file)
            
            # 优先从整体统计信息中获取
            if '__OVERALL_STATS__' in file_stats:
                if 'data_mean' in file_stats['__OVERALL_STATS__'] and 'data_std' in file_stats['__OVERALL_STATS__']:
                    mean = file_stats['__OVERALL_STATS__']['data_mean']
                    std = file_stats['__OVERALL_STATS__']['data_std']
                    print(f"从统计文件获取整体标准化统计信息: mean={mean}, std={std}")
                    return {
                        'mean': [mean],
                        'std': [std]
                    }
            
            # 如果没有整体统计信息，尝试计算所有文件的平均值
            means = []
            stds = []
            for filename, stats in file_stats.items():
                if filename != '__OVERALL_STATS__' and 'data_mean' in stats and 'data_std' in stats:
                    means.append(stats['data_mean'])
                    stds.append(stats['data_std'])
            
            if means and stds:
                avg_mean = sum(means) / len(means)
                avg_std = sum(stds) / len(stds)
                print(f"从统计文件计算平均标准化统计信息: mean={avg_mean}, std={avg_std}")
                return {
                    'mean': [avg_mean],
                    'std': [avg_std]
                }
            
            # 如果没有找到标准化统计信息，返回默认值
            print("统计文件中没有找到标准化统计信息，使用默认值")
            return {
                'mean': [0.0],  # 默认均值
                'std': [1.0]    # 默认标准差
            }
        except Exception as e:
            print(f"从统计文件加载标准化统计信息失败: {str(e)}")
            return {}


# 便捷函数
def get_class_weights_from_stats(stats_file: Optional[str], filenames: Optional[List[str]] = None) -> Optional[List[float]]:
    """
    便捷函数：从统计文件获取类别权重
    
    Args:
        stats_file: 统计文件路径
        filenames: 可选，特定的文件列表
        
    Returns:
        类别权重列表
    """
    return StatsUtil.calculate_class_weights(stats_file, filenames)


def get_file_stages_from_stats(stats_file_path, file_paths, fallback_to_compute=True):
    """
    从统计文件获取文件的阶段数量
    
    参数:
    - stats_file_path: 统计文件路径
    - file_paths: 文件路径列表
    - fallback_to_compute: 如果无法从统计文件获取，是否直接计算
    
    返回:
    - 阶段数量列表，顺序与file_paths对应
    """
    # 初始化阶段数量列表
    stages_count = []
    loaded_count = 0
    
    try:
        # 为每个文件获取阶段数量
        for file_path in file_paths:
            # 使用StatsUtil.get_file_stats方法获取文件统计信息
            file_stats = StatsUtil.get_file_stats(file_path, stats_file_path)
            
            if file_stats and 'stages_count' in file_stats:
                stages = file_stats['stages_count']
                loaded_count += 1
            # 如果无法从统计文件获取且允许fallback，则直接计算
            elif fallback_to_compute:
                stages = compute_file_stages(file_path)
            else:
                stages = 0
            
            stages_count.append(stages)

        
        # 检查是否有有效的阶段数量
        if loaded_count > 0:
            return stages_count
        elif fallback_to_compute:
            # 如果没有有效的阶段数量且允许fallback，则直接计算所有文件
            print("统计文件中没有有效数据，直接计算所有文件的阶段数量")
            return [compute_file_stages(file_path) for file_path in file_paths]
        else:
            return [0] * len(file_paths)
    
    except Exception as e:
        print(f"从统计文件获取阶段数量失败: {str(e)}")
        if fallback_to_compute:
            # 出错时，如果允许fallback，则直接计算
            print("统计文件加载出错，直接计算所有文件的阶段数量")
            return [compute_file_stages(file_path) for file_path in file_paths]
        else:
            return [0] * len(file_paths)


def convert_path_to_absolute(relative_path: str) -> str:
    """
    便捷函数：将相对路径转换为绝对路径
    
    Args:
        relative_path: 相对路径
        
    Returns:
        绝对路径
    """
    return StatsUtil.convert_to_absolute_path(relative_path)


def validate_stats_file(stats_file: str) -> Tuple[bool, str]:
    """
    便捷函数：验证统计文件
    
    Args:
        stats_file: 统计文件路径
        
    Returns:
        (是否有效, 消息)
    """
    return StatsUtil.validate_statistics_file(stats_file)


def compute_file_stages(file_path: str) -> int:
    """
    直接从npz文件计算阶段数量
    
    参数:
    - file_path: npz文件路径
    
    返回:
    - 阶段数量
    """
    try:
        data = np.load(file_path)
        if 'label' in data:
            return len(data['label'])
        elif 'labels' in data:
            return len(data['labels'])
        else:
            return 0
    except (FileNotFoundError, KeyError, IOError) as e:
         print(f"计算文件 {os.path.basename(file_path)} 阶段数量失败: {str(e)}")
         return 0